from langchain.tools import BaseTool
from langchain.chains import create_sql_query_chain
from langchain_community.utilities import SQLDatabase
from langchain.agents import tool
from pydantic import BaseModel,Field,ConfigDict
from typing import Type,Any
from dotenv import load_dotenv
from model.my_chat_model import ChatModel

#定义工具类输入的参数格式
class LoginInput(BaseModel):
    name: str = Field(...,description="用户名")
    password: str = Field(...,description="密码")

#定义工具类
class LoginTool(BaseTool):
    #允许模型接受非标准数据类型
    model_config = ConfigDict(arbitrary_types_allowed=True)
    def __init__(self,**kwargs:Any):
        super().__init__(
            name = "get_login_tool",
            description = "主要用于完成系统的登录功能 ，必须输入的参数是用户名和密码",
            **kwargs
        )
    #指定模型的输入参数必须是LoginInput:
    args_schema : Type[BaseModel] = LoginInput
    #定义工具方法
    def _run(self,name:str,password:str):
        #获取大模型
        chat = ChatModel()
        llm = chat.get_line_model()
        #创建Mysql数据库连接
        db = SQLDatabase.from_uri(
            "mysql+pymysql://root:root@localhost:3306/ai",
            include_tables = ["user_info"]
        )
        #创建链
        chain = create_sql_query_chain(llm,db)
        #提问
        question = f"请根据条件用户名是{name}和密码是{password}查询用户信息"
        sql = chain.invoke({"question":question})
        # print(f"sql:{sql}")
        if "```sql" in sql:
            sql = sql.split("```sql")[1].split("```")[0]
        # print(f" .... :{sql}")
        rs=db.run(sql)
        # print(rs)
        return rs

